package com.zjj.lbw.ai;

import jakarta.annotation.Resource;
import org.springframework.ai.chat.client.ChatClient;
import org.springframework.ai.chat.client.advisor.MessageChatMemoryAdvisor;
import org.springframework.ai.chat.client.advisor.QuestionAnswerAdvisor;
import org.springframework.ai.chat.memory.ChatMemory;
import org.springframework.ai.chat.memory.InMemoryChatMemory;
import org.springframework.ai.chat.messages.Media;
import org.springframework.ai.chat.messages.UserMessage;
import org.springframework.ai.chat.model.ChatModel;
import org.springframework.ai.chat.model.ChatResponse;
import org.springframework.ai.chat.model.StreamingChatModel;
import org.springframework.ai.chat.prompt.Prompt;
import org.springframework.ai.chat.prompt.PromptTemplate;
import org.springframework.ai.document.Document;
import org.springframework.ai.embedding.EmbeddingModel;
import org.springframework.ai.evaluation.EvaluationRequest;
import org.springframework.ai.evaluation.EvaluationResponse;
import org.springframework.ai.evaluation.RelevancyEvaluator;
import org.springframework.ai.model.Content;
import org.springframework.ai.openai.OpenAiChatOptions;
import org.springframework.ai.openai.api.OpenAiApi;
import org.springframework.ai.vectorstore.SearchRequest;
import org.springframework.ai.vectorstore.VectorStore;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.core.io.ClassPathResource;
import org.springframework.util.MimeTypeUtils;
import org.springframework.web.bind.annotation.GetMapping;
import org.springframework.web.bind.annotation.RequestParam;
import org.springframework.web.bind.annotation.RestController;
import org.springframework.web.servlet.mvc.method.annotation.SseEmitter;
import reactor.core.publisher.Flux;

import java.io.IOException;
import java.util.List;
import java.util.Map;

import static org.springframework.ai.chat.client.advisor.AbstractChatMemoryAdvisor.CHAT_MEMORY_CONVERSATION_ID_KEY;

@RestController
public class ChatController {

    @Autowired
    private ChatClient chatClient;

    @Resource
    private ChatModel chatModel;

    @Resource
    private StreamingChatModel streamingChatModel;

    @Resource
    private EmbeddingModel embeddingModel;

    @Autowired
    private DocumentService documentService;

    @Autowired
    private VectorStore vectorStore;

    @GetMapping("/clientChat")
    public String clientChat() {
        return this.chatClient.prompt()
                .user("你是谁")
                .call()
                .content();
    }

    @GetMapping("/streamClientChat")
    public SseEmitter streamClientChat() {
        SseEmitter sseEmitter = new SseEmitter();
        Flux<String> stream = chatClient.prompt()
                .user("你是谁")
                .stream()
                .content();
        stream.subscribe(token -> {
            try {
                sseEmitter.send(token);
            } catch (IOException e) {
                sseEmitter.completeWithError(e);
            }
        }, sseEmitter::completeWithError, sseEmitter::complete);
        return sseEmitter;
    }

    @GetMapping("/systemClientChat")
    public String systemClientChat() {
        return this.chatClient.prompt()
                .system(sp -> sp.param("style", "七言绝句"))
                .user("程序员")
                .call()
                .content();
    }

    @GetMapping("/questionAnswerAdvisor")
    public ChatResponse questionAnswerAdvisor() {
        ChatResponse response = ChatClient.builder(chatModel)
                .build().prompt()
                .advisors(new QuestionAnswerAdvisor(vectorStore, SearchRequest.defaults()))
                .user("apikey是什么")
                .call()
                .chatResponse();

        return response;
    }

    @GetMapping("/chatMemoryChat")
    public String chatMemoryAdvisor(@RequestParam("chatId") String chatId, @RequestParam("message") String message) {
        String content = chatClient
                .prompt()
                .advisors(a -> a.param(CHAT_MEMORY_CONVERSATION_ID_KEY, chatId))
                .user(message)
                .call()
                .content();

        return content;
    }

    @GetMapping("/evaluation")
    public EvaluationResponse evaluation() {
        String userText = "apikey是什么";

        ChatResponse response = ChatClient.builder(chatModel)
                .build().prompt()
                .advisors(new QuestionAnswerAdvisor(vectorStore, SearchRequest.defaults()))
                .user(userText)
                .call()
                .chatResponse();

        var relevancyEvaluator = new RelevancyEvaluator(ChatClient.builder(chatModel));

        EvaluationRequest evaluationRequest = new EvaluationRequest(userText,
                (List<Content>) response.getMetadata().get(QuestionAnswerAdvisor.RETRIEVED_DOCUMENTS), response);

        EvaluationResponse evaluationResponse = relevancyEvaluator.evaluate(evaluationRequest);

        return evaluationResponse;
    }

    @GetMapping("/multimodalChat")
    public String multimodalChat() {
        ClassPathResource imageData = new ClassPathResource("/test.jpg");

        var userMessage = new UserMessage("图片里有什么?",
                List.of(new Media(MimeTypeUtils.IMAGE_PNG, imageData)));

        return chatModel.call(userMessage);
    }

    @GetMapping("/chat")
    public String chat() {
        return chatModel.call("你是谁");
    }

    @GetMapping("/streamChat")
    public SseEmitter streamChat() {
        SseEmitter sseEmitter = new SseEmitter();
        Flux<String> stream = streamingChatModel.stream("你是谁");
        stream.subscribe(token -> {
            try {
                sseEmitter.send(token);
            } catch (IOException e) {
                sseEmitter.completeWithError(e);
            }
        }, sseEmitter::completeWithError, sseEmitter::complete);
        return sseEmitter;
    }

    // Tool
    // Function Call
    @GetMapping("/toolChat")
    public String toolChat() {
        Prompt prompt = new Prompt("今天是几月几号", OpenAiChatOptions.builder().withFunction("dateTool").build());
        return chatModel.call(prompt).getResult().getOutput().getContent();
    }

    @GetMapping("/newsToolChat")
    public String newsToolChat() {
        Prompt prompt = new Prompt("今天有哪些热点新闻", OpenAiChatOptions.builder().withFunction("playwrightTool").build());
        return chatModel.call(prompt).getResult().getOutput().getContent();
    }

    @GetMapping("/embedding")
    public List<Double> embedding(@RequestParam String message) {
        return embeddingModel.embed(message);
    }

    @GetMapping("/documentEmbedding")
    public List<Document> documentEmbedding() {
        return documentService.loadText();
    }

    @GetMapping("/documentSearch")
    public List<Document> documentSearch(@RequestParam String message) {
        return documentService.search(message);
    }

    @GetMapping("/ragChat")
    public String ragChat(@RequestParam String question) {

        // 向量搜索
        List<Document> documentList = documentService.search(question);

        // 提示词模板
        PromptTemplate promptTemplate = new PromptTemplate("{userMessage}\n\n 用以下信息回答问题:\n {contents}");

        // 组装提示词
        Prompt prompt = promptTemplate.create(Map.of("userMessage", question, "contents", documentList));

        // 调用大模型
        return chatModel.call(prompt).getResult().getOutput().getContent();
    }
}
